#include <opencv4/opencv2/opencv.hpp>
#include "acl/acl.h"
#include <time.h>

using namespace cv;
using namespace std;

namespace {
    const float min_chn_0 = 123.675;
    const float min_chn_1 = 116.28;
    const float min_chn_2 = 103.53;
    const float var_reci_chn_0 = 0.0171247538316637;
    const float var_reci_chn_1 = 0.0175070028011204;
    const float var_reci_chn_2 = 0.0174291938997821;
    const int channel = 3;
    const int modelHeight = 256;
    const int modelWeight = 256;
    const int classNum = 30;
}

float * ProcessInput(const string testImgPath)
{
    // 原始图片
    Mat srcImage = imread(testImgPath);

    // resize后的图片
    Mat resizedImage;
    resize(srcImage, resizedImage, Size(256, 256));

    // 分配需要的内存空间
    int input_data_size = channel * modelHeight * modelWeight * sizeof(float);
    float * imageBytes = (float*)malloc(input_data_size);
    memset(imageBytes, 0, input_data_size);

    // 图片标准化
    uint8_t bgrToRgb = 2;
    float meanRgb[3] = {min_chn_2, min_chn_1, min_chn_0};
    float stdRgb[3]  = {var_reci_chn_2, var_reci_chn_1, var_reci_chn_0};
    for (int c = 0; c < channel; ++c)
    {
        for (int h = 0; h < modelHeight; ++h)
        {
            for (int w = 0; w < modelWeight; ++w)
            {
                int dstIdx = (bgrToRgb - c) * modelHeight * modelWeight + h * modelWeight + w;
                imageBytes[dstIdx] =  static_cast<float>((resizedImage.at<cv::Vec3b>(h, w)[c] -
                                                         1.0f*meanRgb[c]) * 1.0f*stdRgb[c] );
            }
        }
    }
    return imageBytes;
}

void printMax(float * out_class_data)
{
    // create map<confidence, class> and sorted by maximum
    map<float, unsigned int, greater<float>> resultMap;
    for (unsigned int j = 0; j < classNum; ++j) {
        resultMap[*out_class_data] = j;
        out_class_data++;
    }

    // 计算softmax的sum
    double totalValue=0.0;
    for (auto it = resultMap.begin(); it != resultMap.end(); ++it) {
        totalValue += exp(it->first);
    }

    // get max <confidence, class>
    float confidence = resultMap.begin()->first;
    unsigned int index = resultMap.begin()->second;
    string line = format("label:%d  conf:%lf ", index, exp(confidence) / totalValue);   
    cout << line << endl;
}


int main()
{
    // ***************************** 初始准备工作 *****************************
    aclInit("");        // 初始化
    aclrtSetDevice(0);  // 设置使用第1块NPU

    // 加载模型
    uint32_t modelId;
    aclError ret = aclmdlLoadFromFile("./resnet50.om", &modelId);



    // ***************************** 准备模型输入 *****************************
    // 对输入图片进行处理
    float * cpu_input_buffer = ProcessInput("./16.jpeg");    // cpu_input_buffer是float32 [1, 3, 256, 256]

    // 将cpu上数据拷贝到NPU上
    void * npu_input_buffer;
    int input_data_size = channel * modelHeight * modelWeight * sizeof(float);
    aclrtMalloc(&npu_input_buffer, input_data_size, ACL_MEM_MALLOC_HUGE_FIRST);         // 在NPU上分配内存
    aclrtMemcpy(npu_input_buffer, input_data_size, cpu_input_buffer, input_data_size, 
                ACL_MEMCPY_DEVICE_TO_DEVICE);                                           // 把CPU上数据拷贝到NPU中
    // 把NPU中分配的显存绑定到aclmdlDataset结构体上
    aclDataBuffer * inputData = aclCreateDataBuffer(npu_input_buffer, input_data_size);
    aclmdlDataset * inputDataset = aclmdlCreateDataset();
    aclmdlAddDatasetBuffer(inputDataset, inputData);



    // ***************************** 准备模型输出空间 *****************************
    // 在npu上分配空间
    void * npu_output_buffer;
    int output_data_size = classNum * sizeof(float);
    aclrtMalloc(&npu_output_buffer, output_data_size, ACL_MEM_MALLOC_HUGE_FIRST);


    // 把NPU中分配的显存绑定到aclmdlDataset结构体上
    aclDataBuffer * outputData = aclCreateDataBuffer(npu_output_buffer, output_data_size);
    aclmdlDataset * outputDataset = aclmdlCreateDataset();
    aclmdlAddDatasetBuffer(outputDataset, outputData);



    // ***************************** 模型执行推理 *****************************
    // ret = aclmdlExecute(modelId, inputDataset, outputDataset);
    
    // stream管理
    aclrtStream * stream;
    aclrtCreateStream(stream);
    aclmdlExecuteAsync(modelId, inputDataset, outputDataset, stream);

    // 增加cpu执行的代码
    
    aclrtSynchronizeStream(stream);



    // ***************************** 模型输出后处理 *****************************
    float * out_class_data = (float*)malloc(output_data_size);
    aclrtMemcpy((void*)out_class_data, output_data_size, npu_output_buffer, output_data_size, ACL_MEMCPY_DEVICE_TO_DEVICE);
    printMax(out_class_data);



    // ***************************** 释放昇腾资源 *****************************
    aclrtFree(npu_input_buffer);
    aclmdlDestroyDataset(inputDataset);
    aclrtFree(npu_output_buffer);
    aclmdlDestroyDataset(outputDataset);
    aclmdlUnload(modelId);
    aclrtResetDevice(0);
    aclFinalize();

    return 0;
}

